import math
from enum import Enum
# from MelodieInfra.
from typing import Callable, List, Optional, Union, TYPE_CHECKING

from RailSim.utils.states import StateMachine, Transition
# from numba.experimental import jitclass
# from numba import int64, float64
if TYPE_CHECKING:
    from ..railroad.railroad import Railroad, Block
from .locomotive import w_i, Locomotive


ctrl_period = 5
dt = float(ctrl_period)


class TrainState(int, Enum):
    STOPPED = 0
    ACCELERATING = 1
    BRAKING = 2
    EVEN_SPEED = 10


class Triggers(int, Enum):
    ACCELERATE = 2
    BRAKE = 3
    BRAKE_TO_STOP = 4
    KEEP_EVEN_SPEED = 10


class TrainFSM(StateMachine):
    def __init__(self, train: "Train"):
        transitions = [
            Transition(Triggers.ACCELERATE,
                       TrainState.STOPPED, TrainState.ACCELERATING),
            Transition(Triggers.ACCELERATE,
                       TrainState.BRAKING, TrainState.ACCELERATING),
            Transition(Triggers.ACCELERATE,
                       TrainState.EVEN_SPEED, TrainState.ACCELERATING),
            Transition(Triggers.ACCELERATE,
                       TrainState.ACCELERATING, TrainState.ACCELERATING),

            Transition(Triggers.KEEP_EVEN_SPEED,
                       TrainState.ACCELERATING, TrainState.EVEN_SPEED),
            Transition(Triggers.KEEP_EVEN_SPEED,
                       TrainState.EVEN_SPEED, TrainState.EVEN_SPEED),
            Transition(Triggers.KEEP_EVEN_SPEED,
                       TrainState.BRAKING, TrainState.EVEN_SPEED),

            Transition(Triggers.BRAKE, TrainState.EVEN_SPEED,
                       TrainState.BRAKING),
            Transition(Triggers.BRAKE, TrainState.BRAKING, TrainState.BRAKING),
            Transition(Triggers.BRAKE_TO_STOP,
                       TrainState.BRAKING, TrainState.STOPPED),
        ]
        initial = TrainState.STOPPED
        self.train = train
        self.trigger_calculators: Callable[[], Triggers] = {  # 根据当前状态，计算触发器
            TrainState.STOPPED: self.on_stopped,
            TrainState.BRAKING: self.on_braking,
            TrainState.ACCELERATING: self.on_accel,
            TrainState.EVEN_SPEED: self.on_even_speed
        }
        super().__init__(transitions, initial)

    def free_blk_speed_map(self, blk_num: int) -> float:
        """前方空闲区段数量与速度的映射"""
        if blk_num == 0:
            return 20.0
        elif blk_num == 1:
            return 35.0
        elif blk_num == 2:
            return 50.0
        else:
            return self.train.max_speed

    def speed_limit_by_distance(self, distance: float) -> float:
        """基于与停车点的距离计算出的限速"""
        if distance > 2000:
            return self.train.max_speed
        elif 1000 < distance <= 2000:
            return 50
        elif 200 < distance <= 1000:
            return 35
        elif 50 < distance < 100:
            return 20
        else:
            return 0

    def calc_max_speed(self):
        """计算列车的速度上限"""
        distance_to_destination = (
            self.train.current_rail.length - self.train.destination.position) - self.train._distance
        free_blocks = self.train.get_free_blocks()
        # print("free_blocks", free_blocks,
        #       'distance-to-destination', distance_to_destination)

        # 前方区间有无车占用导致的速度限制
        speed_limit_by_blocks = self.free_blk_speed_map(free_blocks)
        speed_limit_by_distance = self.speed_limit_by_distance(
            distance_to_destination)
        print('speed-limit', speed_limit_by_distance,
              speed_limit_by_blocks, distance_to_destination)
        return min(speed_limit_by_blocks, speed_limit_by_distance)
        # if distance_to_destination>=2000:
        #     return
        # if distance_to_destination > 50:  # 如果距离停车点超过50米，就返回区间允许的速度
        #     return speed_limit_by_blocks
        # else:
        #     return 0  # 返回0

    def speed_control(self):
        if self.train.max_speed-5 <= self.train.speed <= self.train.max_speed:
            # print("should turn to even speed")
            self.process(Triggers.KEEP_EVEN_SPEED)
        elif self.train.speed > self.train.max_speed:
            self.process(Triggers.BRAKE)
        else:
            self.process(Triggers.ACCELERATE)

    def on_stopped(self):
        if self.calc_max_speed() > 1e-9:
            self.process(Triggers.ACCELERATE)

    def on_braking(self):
        self.train.update_brake()
        speed = self.calc_max_speed()
        print("max-speed when braking", speed, self.train.speed)
        print('bounds', speed-5, self.train.speed, speed)
        if speed-5 <= self.train.speed < speed:
            self.process(Triggers.KEEP_EVEN_SPEED)
        elif self.train.speed < speed-5:
            print('re-accel!')
            self.process(Triggers.ACCELERATE)

    def on_even_speed(self):
        speed = self.calc_max_speed()
        if self.train.speed > speed:
            self.train.target_speed = speed
            # self.train.update_brake()
            self.process(Triggers.BRAKE)
        elif self.train.speed <= speed-5:
            self.train.target_speed = speed
            self.process(Triggers.ACCELERATE)

    def on_accel(self):
        self.train.target_speed = self.calc_max_speed()
        self.train.update_speed()

        self.speed_control()

    def step(self):
        self.trigger_calculators[self.state]()
        # self.train


DIRECTION_UNSET = 0
DIRECTION_DOWN = 1
DIRECTION_UP = 2


class Train:
    loco_props: Locomotive
    _accel: float
    _velocity: float
    _mass: float
    _distance: float
    target_speed: float
    direction: int
    current_block_id: int
    _cars_num: int
    max_speed: float

    def __init__(self):
        self.loco_props = Locomotive()
        self._accel = 0.0  # 当前加速度(m/s^2)
        self._velocity = 0.0  # 当前速度(m/s)
        self._mass = 5000*1000  # 质量(kg)
        self._distance = 0.0
        self.max_speed = 0  # 最大限速
        self.target_speed = 0  # 目标速度，单位为km/h

        self.direction: int = DIRECTION_UP
        self.current_block_id = 0

        # 车厢数量，为列车重量减去车头重量之后，除以60吨/节
        self._cars_num = int(
            (self._mass/1000 - self.loco_props.adhesionMass)/60)+1

        self.stop_at: 'Block' = None
        self.destination: 'Block'
        self.waiting = False

        self.rails: List[Railroad] = []
        self.current_rail: Railroad = None
        self.states = TrainFSM(self)

    def request_wait(self, block: 'Block'):
        """向列车发出调度指令，令其进入最近的会让区间待避"""
        self.stop_at = block

    def cancel_wait(self, rail: "Railroad"):
        """取消列车的待避指令，令列车立即前进"""
        self.stop_at = None
        self.waiting = False
        rail.set_train_running(self, self.current_block_id)

    @property
    def speed(self) -> float:
        return self._velocity*3.6

    def cars_friction(self, v: float) -> float:
        """
        车厢的阻力，单位为kN
        """
        assert v >= -1e-9, v
        return (1.07+0.0011*v+0.000236*v**2+w_i)*(self._mass*9.81/1000)/1000

    def total_friction(self, v: float):
        return self.cars_friction(v)+self.loco_props.friction(v)

    def total_accel(self, v: float):
        # 算出最大的牵引力(kN)
        total_max_traction = self.loco_props.max_traction(
            v) - self.total_friction(v)
        return total_max_traction*1000/self._mass  # 算出加速度

    def increase_speed_simulation(self):
        dt = 1.0  # 时间步为1.0秒
        velocities = []
        accels = []
        for i in range(1000):
            self._accel = self.total_accel(self.speed)
            self._velocity += self._accel*dt
            velocities.append(self.speed)
            accels.append(self._accel)
        return range(1000), velocities, accels

    def update_speed(self):
        self._accel = self.total_accel(self.speed)
        self._velocity += self._accel*dt
        if self._velocity <= 0:
            self._velocity = 0.0

    def update_brake(self):
        self._accel = -0.3
        self._velocity += self._accel*dt
        if self._velocity <= 0:
            self._velocity = 0

    def _update(self, rail: "Railroad"):
        if self.speed < self.target_speed-5:
            self.update_speed()
        elif self.target_speed > 0 and self.target_speed-5 <= self.speed <= self.target_speed:
            pass
        else:
            self.update_brake()
        self._distance += self._velocity*dt

    def get_free_blocks(self):
        if self.direction == DIRECTION_UP:
            return self.current_rail.free_blocks_down(self.current_block_id)
        else:
            return self.current_rail.free_blocks_down(self.current_block_id)

    def step_down(self, step: int, rail: 'Railroad'):
        current_block = rail.blocks[self.current_block_id]
        if self._distance > current_block.end:
            if self.current_block_id < len(rail.blocks)-1:
                last_block = self.current_block_id
                self.current_block_id += 1
                rail.move_train_block(last_block, self.current_block_id)

        if step % 10 == 0:  # 司机每一段时间都判断一下前方是否无车
            free_blocks = rail.free_blocks_up(self.current_block_id)
            if free_blocks >= 2:
                self.target_speed = self.max_speed
            elif free_blocks == 1:
                self.target_speed = 20
            else:
                self.target_speed = 0

    def step_up(self, step: int):
        rail = self.current_rail
        current_block = rail.blocks[self.current_block_id]
        if (rail.length - self._distance) < current_block.start:
            if self.current_block_id >= 0:
                last_block = self.current_block_id
                self.current_block_id -= 1
                # print(last_block, self.current_block_id, rail.train_occupations)
                rail.move_train_block(last_block, self.current_block_id)

        # if step % 10 == 0:  # 司机每一段时间都判断一下前方是否无车，根据这个标准来调整速度
            # speed = self.max_speed
            # if self.stop_at is not None:
            #     should_stop_at = (rail.length - self.stop_at.position) - \
            #         self._distance

            #     if 1000 <= should_stop_at < 2000:
            #         speed = 60
            #         # return
            #     elif 200 <= should_stop_at < 400:
            #         speed = 30
            #         speed = 0
            #     elif should_stop_at < 200:
            #         speed = 0
            #         self.waiting = True
            #         self._velocity = 0
            #         self._accel = 0
            #         rail.set_train_waiting(self.current_block_id)
            #     print("should stop at:", should_stop_at,
            #           (rail.length - self.stop_at.position)/1000.0)
            # free_blocks = rail.free_blocks_down(self.current_block_id)
            # if free_blocks >= 2:
            #     self.target_speed = speed
            # elif free_blocks == 1:
            #     self.target_speed = 20
            # else:
            #     self.target_speed = 0
            # print('target_speed', self.target_speed)

    def init_status(self, rail: "Railroad", direction: int):
        self.direction = direction

        self.current_block_id = rail.put_train(self, self.direction)

    def update(self, step: int, rail: 'Railroad'):
        if self.waiting:
            return
        if self.direction == DIRECTION_UP:
            self.step_up(step, rail)
        else:
            self.step_down(step, rail)
        self._update(rail)

    def step(self, step: int):
        if step % ctrl_period == 0:
            self.states.step()
        self._distance += self._velocity*1
        self.step_up(step)
